
from Crypto.PublicKey import RSA
from Crypto.Cipher import ARC4

from multitask import FDAction, taskmgr, sleep, Queue, Timeout, Return

import threading, re, pprint, urlparse
import socket, sys, os, errno
import pickle, hashlib, binascii


IOBLOCK = (errno.EWOULDBLOCK, errno.EINPROGRESS, errno.ENOBUFS, errno.EISCONN)

class DNSResolver(threading.Thread):
    Event = threading.Event()
    Request = set()
    Result = {}
    def __init__(self, name, daemon):
        threading.Thread.__init__(self)
        self.name = name
        self.daemon = daemon

    def run(self):
        while True:
            self.Event.wait()
            while self.Request:
                addr = self.Request.pop()
                if addr not in self.Result:
                    try: info = socket.getaddrinfo(addr[0], addr[1], socket.AF_INET)
                    except socket.gaierror:
                        try: info = socket.getaddrinfo(addr[0], addr[1], socket.AF_INET6)
                        except socket.gaierror as e: info = e
                    self.Result[addr] = info
                print self.Result
            self.Event.clear()
            
    @classmethod
    def request(cls, addr):
        cls.Request.add(addr)
        cls.Event.set()
        
    @classmethod
    def query(cls, addr):
        info = cls.Result.get(addr)
        if isinstance(info, socket.gaierror): raise info
        return info        

DNSResolver(name="dnsd", daemon=True).start()

def connect(*args):
    info = DNSResolver.query(args)
    while info is None:
        DNSResolver.request(args)
        yield sleep(0.05)
        info = DNSResolver.query(args)
    sock = socket.socket(*info[0][:3])
    sock.setblocking(0)
    code = sock.connect_ex(info[0][-1])
    if code not in IOBLOCK:
        raise socket.error(code)
    raise Return(sock)

def sendall(sock, data):
    def _sendall():
        buf = _sendall._buf
        try:
            buf = buf[sock.send(buf):]
            if buf: return Ellipsis
        except socket.error, e:
            if e.errno in IOBLOCK:
                return Ellipsis
            raise
        finally:
            _sendall._buf = buf
    _sendall._buf = data
    return FDAction(sock, _sendall, write=True)

def recvall(sock, until):
    if isinstance(until, int):
        if until == 0: return ""
        def _recvall():
            data = sock.recv(16384 if until < 0 else min(until - _recvall._pos, 16384))
            if not data:
                result = "".join(_recvall._buf)
                if until < 0 or _recvall._pos == until:
                    return result
                raise EOFError(result)
            _recvall._buf.append(data)
            _recvall._pos += len(data)
            if _recvall._pos == until:
                return "".join(_recvall._buf)
            return Ellipsis
    elif isinstance(until, str):
        def _recvall():
            buf, pos = _recvall._buf, _recvall._pos
            timeout = sock.gettimeout()
            sock.settimeout(0)
            try:
                all = len(until)
                while pos != all:
                    ch = sock.recv(1)
                    if not ch:
                        raise EOFError("".join(buf))
                    buf.append(ch)
                    pos = pos + 1 if until[pos] == ch else 0
                return "".join(buf)
            except socket.error, e:
                if e.errno in (errno.EWOULDBLOCK,):
                    return Ellipsis
                raise
            finally:
                _recvall._pos = pos
                sock.settimeout(timeout)
    _recvall._buf = []
    _recvall._pos = 0
    return FDAction(sock, _recvall, read=True)

def trans_packit(request, PICKLE):
    if PICKLE: # str->list
        if not isinstance(request, str):
            return pickle.dumps([list(request)], pickle.HIGHEST_PROTOCOL)
        else:
            raise NotImplemented
    else: # list->str
        if isinstance(request, str):
            return request
        else:
            method, path, headers, payload = request
            packet = [method, " ", path, " HTTP/1.1\r\n"]
            for header in headers:
                packet.append("%s: %s\r\n" % header)
            packet.extend(["\r\n", payload])
            return "".join(packet)

def trans_unpack(response, PICKLE):
    if PICKLE:
        return pickle.loads(response)[0]
    else:
        return response

def GAEClient(url, key="proxy.key", wrap=""):
    print "new GAEClient Instance start: url= " + url
    cipher = RSA.importKey(open(key).read())
    PICKLE = False
    rs = urlparse.urlsplit(url)
    template = ("POST %s%%s HTTP/1.0\r\n" % rs.path +
                "Host: %s\r\n" % rs.netloc +
                "Content-Type: application/octet-stream\r\n" +
                "Connection: close\r\n" +
                "X-Authorization: %s\r\n" +
                "Content-Length: %s\r\n\r\n")
    while True:
        try:
            request, caller = yield GAERequest.get(600)
            yield caller.put(None)
        except Timeout: return
        session = binascii.b2a_hex(os.urandom(32))[:64]
        request = trans_packit(request, PICKLE=PICKLE)
        trans = ARC4.new(session).encrypt(request)
        token = cipher.encrypt(session + hashlib.sha512(trans).digest(), 0)[0]
        token = binascii.b2a_base64(token).strip()
        packet = template % ("", token, len(trans))

        sock = yield connect(rs.hostname, rs.port or 80)
        yield sendall(sock, packet)
        yield sendall(sock, trans)
        trans = yield recvall(sock, -1)

        packet, trans = trans.split("\r\n\r\n", 1)
        #print packet
        token = re.findall("X-Validation:(.*)", packet, re.I)[0]
        token = binascii.a2b_base64(token)
        #print "====" * 5, hashlib.sha512(request).digest() == token[0:64],
        #print hashlib.sha512(trans).digest() == token[64:128], "====" * 5
        response = ARC4.new(session).decrypt(trans)
        response = trans_unpack(response, PICKLE)
        yield caller.put(response)
        #print pprint.pformat(pickle.loads(response)) if PICKLE else response[:1000]

GAERequest = Queue()
GAEInfo = dict(url="http://lyricconch-proxy.appspot.com/fetch.py",
               key="proxy.key",
               wrap="")
def gae_fetch(request):
    caller = Queue(maxsize=1)
    yield GAERequest.put((request, caller))
    while True:
        try: yield caller.get(timeout=1); break
        except Timeout:
            taskmgr.add(GAEClient(**GAEInfo))
    raise Return((yield caller.get()))

if __name__ == "__main__":

    def a(addr):
        print addr
        sock = yield connect(addr, 80)
        print sock.getpeername()

    def b(text):
        for dummy in xrange(100):
            print (yield text),
            yield sleep(1)

    taskmgr.add(b("z"))
    taskmgr.add(a("www.google.com"))
    taskmgr.add(a("www.mozilla.org"))
    taskmgr.add(a("www.google.com"))

    def c(i):
        print "c #%s start" % i
        request = ("GET", "http://www.google.com/", {}, "")
        caller = Queue(maxsize=1)
        yield GAERequest.put((request, caller))
        while True:
            try: (yield caller.get(timeout=1)); break
            except Timeout:
                print "Timeout"
                taskmgr.add(GAEClient(**GAEInfo))
        print "c #%s got" % i , (yield caller.get()).split("\n", 1)[0]

    def d():
        yield taskmgr.add(c(-1))
        for i in xrange(10):
            yield sleep(0.5)
            taskmgr.add(c(i))

    taskmgr.add(d())
    try:
        taskmgr.run()
    except:
        sys.stderr.flush()

